package org.example.algorithm.unionfind;

public class UnionFind {
    private int[] parent;
    private int[] rank;

    public UnionFind(int size) {
        parent = new int[size];
        rank = new int[size];
        for (int i=0;i<size;i++) {
            parent[i] = i;
            rank[i] = 1;
        }
    }

    public int find(int x) {
        if (parent[x] != x) {
            parent[x] = find(parent[x]);
        }
        return parent[x];
    }

    // 合并操作（按秩合并）
    public void union(int x, int y) {
        int rootX = find(x);
        int rootY = find(y);
        if (rootX != rootY) {
            if (rank[rootX] > rank[rootY]) {
                parent[rootY] = rootX;
            } else if(rank[rootX] < rank[rootY]) {
                parent[rootX] = rootY;
            } else {
                parent[rootY] = rootX;
                rank[rootX]++;
            }
        }
    }

    // 判断两个元素是否连通
    public boolean isConnected(int x, int y) {
        return find(x) == find(y);
    }

    public static void main(String[] args) {
        UnionFind uf = new UnionFind(10); // 初始化包含10个元素的并查集

        // 合并一些元素
        uf.union(1, 2);
        uf.union(2, 3);
        uf.union(4, 5);

        // 检查连通性
        System.out.println(uf.isConnected(1, 3)); // true
        System.out.println(uf.isConnected(1, 4)); // false

        // 合并更多元素
        uf.union(3, 4);
        System.out.println(uf.isConnected(1, 5)); // true

    }
}
